import random
import numpy as np
import scipy.special as sc
from chaco.api import *
from traits.api import *
from traitsui.api import *
from chaco import default_colormaps
from enable.api import ComponentEditor
from chaco.tools.api import LineInspector, PanTool, ZoomTool
colormaps = list(default_colormaps.color_map_name_dict.keys())
for boring in "bone gray yarg gist_gray gist_yarg Greys".split():
    colormaps.remove(boring)
class Absorber(HasTraits):
    epsilon=Complex(4-1j)
    mu=Complex(2-1j)
    thickness=Float(2)
    traits_view=View(HGroup(
        Item('epsilon',label='介电常数'),
        Item('mu',label='磁导率'),
        show_border=True,
    ))
table_editor = TableEditor(
    columns=[
        ObjectColumn(name='epsilon',label='介电常数',width=.3),
        ObjectColumn(name='mu',label='磁导率',width=.3),
        ObjectColumn(name='thickness',label='厚度/mm',width=.3),
    ],
    deletable=True,
    sortable=False,
    reorderable=True,
    show_toolbar=True,
    row_factory=Absorber,
)
class Form(HasTraits):
    options=Enum("TE S11","TE S12",'TE S21','TE S22',"TM S11","TM S12",'TM S21','TM S22',)
    min_x = CFloat(0)
    max_x = CFloat(90)
    npts_x = CInt(91)
    title_x='bistatic angle / deg'
    min_y = CFloat(2)
    max_y = CFloat(20)
    npts_y = CInt(181)
    title_y='Frequency / GHz'
    value_label='Incident Angle / deg'
    xs = Array
    ys = Array
    zs = Dict
    plot_x=Instance(Plot)
    plot_y=Instance(Plot)
    container = Instance(HPlotContainer)
    polyplot = Instance(ContourPolyPlot)
    lineplot = Instance(ContourLinePlot)
    cross_plot_x = Instance(Plot)
    cross_plot_y = Instance(Plot)
    colorbar = Instance(ColorBar)
    num_levels = Int(15)
    colormap = Enum(colormaps)
    _image_index = Instance(GridDataSource)
    _image_value = Instance(ImageData)
    _cmap = Callable(default_colormaps.viridis)
    edit_model = Button
    xMaxIndex=Property(Int,observe=['xs'])
    yMaxIndex=Property(Int,observe=['ys'])
    xIndex=Int
    yIndex=Int

    background=Instance(Absorber,())
    backgroundType=Enum('自由空间','导体','介质')
    absorbers = List(Absorber,[
        Absorber(epsilon=11-0.4j,mu=1.1-0.9j),
    ])
    def _get_xMaxIndex(self):
        return np.array(self.xs).size-1
    def _get_yMaxIndex(self):
        return np.array(self.ys).size-1
    def traits_init(self):
        self.compute_model()
        self.create_plot()
        self.create_plotxy()
        self.update()
    def compute_model(self):#这里只计算了另一侧为自由空间时的情形。还要考虑金属背板、材料衬底的情形。
        self.xs = np.linspace(self.min_x, self.max_x, self.npts_x)
        self.ys = np.linspace(self.min_y, self.max_y, self.npts_y)
        f,phi=np.meshgrid(self.ys,self.xs,indexing='ij')
        sinphi=sc.sindg(phi)
        cosphi=sc.cosdg(phi)
        sinphi2=sinphi**2
        mats=[]
        mat0=np.zeros((self.npts_y,self.npts_x,2,2),dtype=complex)
        mat0[:,:,0,0]=np.sqrt(cosphi)
        mat0[:,:,1,1]=1/np.sqrt(cosphi)
        mats.append(mat0)
        for absorber in self.absorbers:
            e=absorber.epsilon
            m=absorber.mu
            d=absorber.thickness
            k=np.sqrt(e*m-sinphi2)
            kd=2*np.pi*f*d/300
            coskd=np.cos(kd*k)
            sinkd=np.sin(kd*k)
            mat=np.zeros_like(mat0,dtype=complex)
            mat[:,:,0,0]=coskd
            mat[:,:,1,1]=coskd
            mat[:,:,0,1]=1j*sinkd*m/k
            mat[:,:,1,0]=1j*sinkd*k/m
            mats.append(mat)
        mat0=np.zeros((self.npts_y,self.npts_x,2,2),dtype=complex)
        mat0[:,:,0,0]=1/np.sqrt(cosphi)
        mat0[:,:,1,1]=np.sqrt(cosphi)
        mats.append(mat0)
        abcd=np.zeros_like(mat0,dtype=complex)
        abcd[:, :, 0, 0] = 1
        abcd[:, :, 1, 1] = 1
        for mat in mats:
            abcd=abcd@mat
        s=np.zeros_like(mat0,dtype=complex)
        a=abcd[:,:,0,0]
        b=abcd[:,:,0,1]
        c=abcd[:,:,1,0]
        d=abcd[:,:,1,1]
        denominator=a+b+c+d
        s[:,:,0,0]=a+b-c-d
        s[:,:,0,1]=2*(a*d-b*c)
        s[:,:,1,0]=2
        s[:,:,1,1]=-a+b-c+d
        s/=denominator[:,:,np.newaxis,np.newaxis]
        self.zs['TE S11']=20*np.log10(np.abs(s[:,:,0,0]))
        self.zs['TE S12']=20*np.log10(np.abs(s[:,:,0,1]))
        self.zs['TE S21']=20*np.log10(np.abs(s[:,:,1,0]))
        self.zs['TE S22']=20*np.log10(np.abs(s[:,:,1,1]))
        mats=[]
        mat0=np.zeros((self.npts_y,self.npts_x,2,2),dtype=complex)
        mat0[:,:,0,0]=np.sqrt(cosphi)
        mat0[:,:,1,1]=1/np.sqrt(cosphi)
        mats.append(mat0)
        for absorber in self.absorbers:
            e=absorber.epsilon
            m=absorber.mu
            d=absorber.thickness
            k=np.sqrt(e*m-sinphi2)
            kd=2*np.pi*f*d/300
            coskd=np.cos(kd*k)
            sinkd=np.sin(kd*k)
            mat=np.zeros_like(mat0,dtype=complex)
            mat[:,:,0,0]=coskd
            mat[:,:,1,1]=coskd
            mat[:,:,0,1]=1j*sinkd*e/k
            mat[:,:,1,0]=1j*sinkd*k/e
            mats.append(mat)
        mat0=np.zeros((self.npts_y,self.npts_x,2,2),dtype=complex)
        mat0[:,:,0,0]=1/np.sqrt(cosphi)
        mat0[:,:,1,1]=np.sqrt(cosphi)
        mats.append(mat0)
        abcd=np.zeros_like(mat0,dtype=complex)
        abcd[:, :, 0, 0] = 1
        abcd[:, :, 1, 1] = 1
        for mat in mats:
            abcd=abcd@mat
        s=np.zeros_like(mat0,dtype=complex)
        a=abcd[:,:,0,0]
        b=abcd[:,:,0,1]
        c=abcd[:,:,1,0]
        d=abcd[:,:,1,1]
        denominator=a+b+c+d
        s[:,:,0,0]=a+b-c-d
        s[:,:,0,1]=2*(a*d-b*c)
        s[:,:,1,0]=2
        s[:,:,1,1]=-a+b-c+d
        s/=denominator[:,:,np.newaxis,np.newaxis]
        self.zs['TM S11']=20*np.log10(np.abs(s[:,:,0,0]))
        self.zs['TM S12']=20*np.log10(np.abs(s[:,:,0,1]))
        self.zs['TM S21']=20*np.log10(np.abs(s[:,:,1,0]))
        self.zs['TM S22']=20*np.log10(np.abs(s[:,:,1,1]))
    def create_plotxy(self):
        self.plot_x=Plot(ArrayPlotData())
        self.plot_y=Plot(ArrayPlotData())
        self.updatexy()
        self.plot_x.plot(('index','tes11'), name='TE S11',color=palette11[0])
        self.plot_x.plot(('index','tes12'), name='TE S12',color=palette11[1])
        self.plot_x.plot(('index','tes21'), name='TE S21',color=palette11[2])
        self.plot_x.plot(('index','tes22'), name='TE S22',color=palette11[3])
        self.plot_x.plot(('index','tms11'), name='TM S11',color=palette11[4])
        self.plot_x.plot(('index','tms12'), name='TM S12',color=palette11[5])
        self.plot_x.plot(('index','tms21'), name='TM S21',color=palette11[6])
        self.plot_x.plot(('index','tms22'), name='TM S22',color=palette11[7])
        self.plot_x.x_axis.title = "incident Angle / deg"
        self.plot_x.y_axis.title = "S parameters / dB"
        legend = Legend()
        legend.plots = self.plot_x.plots
        self.plot_x.overlays.append(legend)
        self.plot_y.plot(('index','tes11'), name='TE S11',color=palette11[0])
        self.plot_y.plot(('index','tes12'), name='TE S12',color=palette11[1])
        self.plot_y.plot(('index','tes21'), name='TE S21',color=palette11[2])
        self.plot_y.plot(('index','tes22'), name='TE S22',color=palette11[3])
        self.plot_y.plot(('index','tms11'), name='TM S11',color=palette11[4])
        self.plot_y.plot(('index','tms12'), name='TM S12',color=palette11[5])
        self.plot_y.plot(('index','tms21'), name='TM S21',color=palette11[6])
        self.plot_y.plot(('index','tms22'), name='TM S22',color=palette11[7])
        self.plot_y.x_axis.title = "Frequency / GHz"
        self.plot_y.y_axis.title = "S parameters / dB"
        legend = Legend()
        legend.plots = self.plot_y.plots
        self.plot_y.overlays.append(legend)
    def create_plot(self):
        self._image_index = GridDataSource(np.array([]), np.array([]), sort_order=("ascending", "ascending"))
        image_index_range = DataRange2D(self._image_index)
        self._image_index.observe(self._metadata_changed, "metadata_changed")
        self._image_value = ImageData(data=np.array([]), value_depth=1)
        image_value_range = DataRange1D(self._image_value)
        self.polyplot = ContourPolyPlot(
            index=self._image_index,
            value=self._image_value,
            index_mapper=GridMapper(range=image_index_range),
            levels=self.num_levels,
            color_mapper=self._cmap(image_value_range),
        )
        self.lineplot = ContourLinePlot(
            index=self._image_index,
            value=self._image_value,
            index_mapper=GridMapper(range=self.polyplot.index_mapper.range),
            levels=self.num_levels,
        )
        self.left = PlotAxis(
            orientation="left",
            title=self.title_y,
            mapper=self.polyplot.index_mapper._ymapper,
            component=self.polyplot,
        )
        self.polyplot.overlays.append(self.left)
        self.bottom = PlotAxis(
            orientation="bottom",
            title=self.title_x,
            mapper=self.polyplot.index_mapper._xmapper,
            component=self.polyplot,
        )
        self.polyplot.overlays.append(self.bottom)
        self.polyplot.tools.append(
            PanTool(self.polyplot, constrain_key="shift")
        )
        self.polyplot.overlays.append(
            ZoomTool(component=self.polyplot, tool_mode="box", always_on=False)
        )
        self.polyplot.overlays.append(
            LineInspector(
                component=self.polyplot,
                axis="index_x",
                inspect_mode="indexed",
                write_metadata=True,
                color="white",
                is_listener=True,
            )
        )
        self.polyplot.overlays.append(
            LineInspector(
                component=self.polyplot,
                axis="index_y",
                inspect_mode="indexed",
                write_metadata=True,
                color="white",
                is_listener=True,
            )
        )
        contour_container = OverlayPlotContainer(padding=20, use_backbuffer=True, unified_draw=True)
        contour_container.add(self.polyplot)
        contour_container.add(self.lineplot)
        cbar_index_mapper = LinearMapper(range=image_value_range)
        self.colorbar = ColorBar(
            index_mapper=cbar_index_mapper,
            plot=self.polyplot,
            padding_top=self.polyplot.padding_top,
            padding_bottom=self.polyplot.padding_bottom,
            padding_right=40,
            resizable="v",
            width=30,
        )
        self.cross_plot_x = Plot(
            ArrayPlotData(
                line_index=np.array([]),
                line_value=np.array([]),
                scatter_index=np.array([]),
                scatter_value=np.array([]),
                scatter_color=np.array([]),
            ),
            resizable="h",
            height = 140,
            padding = 20,
        )
        self.cross_plot_x.plot(("line_index", "line_value"), line_style="dot")
        self.cross_plot_x.plot(
            ("scatter_index", "scatter_value", "scatter_color"),
            type="cmap_scatter",
            name="dot",
            color_mapper=self._cmap(image_value_range),
            marker="circle",
            marker_size=8,
        )
        self.cross_plot_x.index_range = self.polyplot.index_range.x_range
        self.cross_plot_y = Plot(
            ArrayPlotData(
                line_index=np.array([]),
                line_value=np.array([]),
                scatter_index=np.array([]),
                scatter_value=np.array([]),
                scatter_color=np.array([]),
            ),
            width=140,
            resizable="v",
            padding=20,
            orientation="v",
            padding_bottom=200,
        )
        self.cross_plot_y.plot(("line_index", "line_value"), line_style="dot")
        self.cross_plot_y.plot(
            ("scatter_index", "scatter_value", "scatter_color"),
            type="cmap_scatter",
            name="dot",
            color_mapper=self._cmap(image_value_range),
            marker="circle",
            marker_size=8,
        )
        self.cross_plot_y.index_range = self.polyplot.index_range.y_range
        self.container = HPlotContainer(
            padding=40,
            fill_padding=True,
            bgcolor="white",
            use_backbuffer=False,
        )
        inner_cont = VPlotContainer(padding=0, use_backbuffer=True)
        inner_cont.add(self.cross_plot_x)
        inner_cont.add(contour_container)
        self.container.add(self.colorbar)
        self.container.add(inner_cont)
        self.container.add(self.cross_plot_y)
    @observe('options')
    def update(self,event=None):
        self.colorbar.index_mapper.range.low = np.nanmin(self.zs[self.options])
        self.colorbar.index_mapper.range.high = np.nanmax(self.zs[self.options])
        self._image_index.set_data(self.xs, self.ys)
        self._image_value.data = self.zs[self.options]
        self.cross_plot_x.data.update_data(line_index_x=self.xs)
        self.cross_plot_y.data.update_data(line_index_y=self.ys)
    def _metadata_changed(self,event=None):
        if "selections" in self._image_index.metadata and len(self._image_index.metadata['selections'])==2:
            x_ndx, y_ndx = self._image_index.metadata["selections"]
            if y_ndx and x_ndx:
                xdata, ydata = self._image_index.get_data()
                xdata, ydata = xdata.get_data(), ydata.get_data()
                self.cross_plot_x.data.update_data(
                    line_index=self.xs,
                    line_value=self._image_value.data[y_ndx, :],
                    scatter_index=np.array([xdata[x_ndx]]),
                    scatter_value=np.array([self._image_value.data[y_ndx, x_ndx]]),
                    scatter_color=np.array([self._image_value.data[y_ndx, x_ndx]]),
                )
                self.cross_plot_y.data.update_data(
                    line_index=self.ys,
                    line_value=self._image_value.data[:, x_ndx],
                    scatter_index=np.array([ydata[y_ndx]]),
                    scatter_value=np.array([self._image_value.data[y_ndx, x_ndx]]),
                    scatter_color=np.array([self._image_value.data[y_ndx, x_ndx]]),
                )
        else:
            self.cross_plot_x.data.update_data(
                line_value=np.array([]),
                scatter_value=np.array([]),
                )
            self.cross_plot_y.data.update_data(
                line_value=np.array([]),
                scatter_value=np.array([]),
                )
    def _colormap_changed(self):
        self._cmap = default_colormaps.color_map_name_dict[self.colormap]
        if self.polyplot is not None:
            value_range = self.polyplot.color_mapper.range
            self.polyplot.color_mapper = self._cmap(value_range)
            value_range = self.cross_plot_x.color_mapper.range
            self.cross_plot_x.color_mapper = self._cmap(value_range)
            self.cross_plot_y.color_mapper = self._cmap(value_range)
        self.container.request_redraw()
    def _num_levels_changed(self):
        if self.num_levels > 3:
            if self.polyplot is not None:
                self.polyplot.levels = self.num_levels
            if self.lineplot is not None:
                self.lineplot.levels = self.num_levels
    def _edit_model_fired(self):
        self.configure_traits(view='model_view')
    class ModelHandler(Handler):
        def closed(self, info, is_ok):
            if is_ok:
                obj = info.object
                obj.compute_model()
                obj.update()
                obj.updatexy()
            else:
                print('Cancel')
    model_view = View(
        Group(
            HGroup(
                Item("npts_x", label="入射角点数"),
                Item("npts_y", label="频率点数"),
            ),
            HGroup(
                Item("min_x", label="最小入射角/deg"),
                Item("min_y", label="最小频率/GHz"),
            ),
            HGroup(
                Item("max_x", label="最大入射角/deg"),
                Item("max_y", label="最大频率/GHz"),
            ),
            HGroup(Item('backgroundType', label='底板类型', style='custom'),
                   Item('background', enabled_when='backgroundType=="介质"', label='底板', style='custom'),
                   show_border=True
                   ),
            Item('absorbers', editor=table_editor, height=200, label='吸波材料'),

        ),
        buttons=["OK", "Cancel"],
        handler=ModelHandler,
    )
    @observe('xIndex,yIndex')
    def updatexy(self,event=None):
        self.plot_x.data.update_data(
            index=self.xs,
            tes11=self.zs['TE S11'][self.yIndex,:],
            tes12=self.zs['TE S12'][self.yIndex,:],
            tes21=self.zs['TE S21'][self.yIndex,:],
            tes22=self.zs['TE S22'][self.yIndex,:],
            tms11=self.zs['TM S11'][self.yIndex,:],
            tms12=self.zs['TM S12'][self.yIndex,:],
            tms21=self.zs['TM S21'][self.yIndex,:],
            tms22=self.zs['TM S22'][self.yIndex,:],
        )
        self.plot_y.data.update_data(
            index=self.ys,
            tes11=self.zs['TE S11'][:,self.xIndex],
            tes12=self.zs['TE S12'][:,self.xIndex],
            tes21=self.zs['TE S21'][:,self.xIndex],
            tes22=self.zs['TE S22'][:,self.xIndex],
            tms11=self.zs['TM S11'][:,self.xIndex],
            tms12=self.zs['TM S12'][:,self.xIndex],
            tms21=self.zs['TM S21'][:,self.xIndex],
            tms22=self.zs['TM S22'][:,self.xIndex],
        )
        self.plot_x.title = f'Frequency: {self.ys[self.yIndex]} GHz'
        self.plot_y.title = f'Incident Angle: {self.xs[self.xIndex]} deg'
    traits_view = View(
        Group(
            Tabbed(
            Group(
                UItem("container", editor=ComponentEditor(size=(800, 600))),
                HGroup(
                    Item("num_levels",label='等高线数量'),
                    Item("colormap",label='颜色'),
                    Item("options",label='选项',style='custom'),
                    spring,
                ),
                label='频率和入射角'
            ),
            Group(
                Group(UItem("plot_x", editor=ComponentEditor(size=(800, 600)))),
                Item('yIndex', editor=RangeEditor(
                    low=0,
                    high_name='yMaxIndex',
                    mode='slider')),
                label='入射角'
            ),
            Group(
                Group(UItem("plot_y", editor=ComponentEditor(size=(800, 600)))),
                Item('xIndex', editor=RangeEditor(
                    low=0,
                    high_name='xMaxIndex',
                    mode='slider'),
                     ),
                label='频率'
            ),
        ),
        HGroup(
            UItem("edit_model",label='设置参数'),
            spring,
        )
        ),
        title="多层材料S参数反射",
        resizable=True,
    )
if __name__ == "__main__":
    form = Form()
    form.configure_traits()
